package BinaryTree
import(
	"fmt"
)

func Abs(a,b int)int{
	if a > b{
		return a - b
	}
	return b - a
}
func Max(a,b int)int{
	if a > b{
		return a
	}else{
		return b
	}
}
type AVLTree struct{
	root	*tree_node
}

type tree_node struct{
	data		int
	parent		*tree_node
	leftChild	*tree_node
	rightChild	*tree_node
}

func New_AVLTree()	*AVLTree{
	tree := &AVLTree{}
	return tree
}

func (tree *AVLTree)Visit(){
	if tree.root == nil{
		fmt.Println("tree empty")
	}
	tree.root.visit()
	fmt.Println()
}

func (tree *AVLTree)InsertData(data int){
	node := &tree_node{data:data}
	tree.InsertNode(node)
}

func (tree *AVLTree)InsertNode(node *tree_node){
	if node == nil{return}

	if tree.root == nil{
		tree.root = node
		tree.root.rightChild = nil
		tree.root.leftChild = nil
		tree.root.parent = nil	
	}else{
		tree.root.insert(node,tree)
	}
}

func (this *tree_node)insert(node *tree_node,tree *AVLTree){
	if node.data < this.data {
		this.insertLeft(node,tree)
	}else if node.data > this.data {
		this.insertRight(node,tree)
	}
	lh,rh := this.getLeftHeight(),this.getRightHeight()
	fmt.Printf("node:%d left height:%d right height:%d\n",this.data,lh,rh)

	if Abs(lh,rh) >= 2 {

		if lh > rh{
			//左边超过高度差
			lch,rch := this.leftChild.getLeftHeight(),this.leftChild.getRightHeight()
			if lch > rch{
				//LL旋转
				this.LLRotate(tree)
			}else{
				this.LRRotate(tree)
			}
		}else{
			lch,rch := this.rightChild.getLeftHeight(),this.rightChild.getRightHeight()
			if rch > lch{
				//RR旋转
				this.RRRotate(tree)
			}else{
				this.rightChild.RLRotate(tree)
			}
		}
	
	}

}

func (this *tree_node)insertLeft(node *tree_node,tree *AVLTree){
	if this.leftChild == nil{
		this.leftChild = node
		node.parent = this
	}else{
		this.leftChild.insert(node,tree)
	}
}

func (this *tree_node)insertRight(node *tree_node,tree *AVLTree){
	if this.rightChild == nil{
		this.rightChild = node
		node.parent = this
	}else{
		this.rightChild.insert(node,tree)
	}
}

func (this *tree_node)LLRotate(tree *AVLTree){
	fmt.Println("LLRotate")
	//左孩子替代自己
	replaceNode := this.leftChild
	//左孩子的右孩子成为自己的左孩子
	if replaceNode.rightChild != nil{
		this.leftChild = replaceNode.rightChild
		replaceNode.rightChild.parent = this
	}else{
		this.leftChild = nil
	}
	//自己成为原本左孩子的右孩子
	if this.parent == nil{
		//root
		tree.root = replaceNode
		replaceNode.parent = nil
	}else{
		replaceNode.parent = this.parent
		this.parent.rightChild = replaceNode
	}
	replaceNode.rightChild = this
	this.parent = replaceNode
}

func (this *tree_node)RRRotate(tree *AVLTree){
	fmt.Println("RRRoate")
	replaceNode := this.rightChild
	
	if replaceNode.leftChild != nil{
		this.rightChild = replaceNode.leftChild
		replaceNode.leftChild.parent = this
	}else{
		this.rightChild = nil
	}

	if this.parent == nil{
		tree.root = replaceNode
		replaceNode.parent = nil
	}else{
		replaceNode.parent = this.parent
		this.parent.leftChild = replaceNode
	}
	replaceNode.leftChild = this
	this.parent = replaceNode
}

func (this *tree_node)LRRotate(tree *AVLTree){
	this.leftChild.RRRotate(tree)
	this.LLRotate(tree)
}

func (this *tree_node)RLRotate(tree *AVLTree){
	this.rightChild.LLRotate(tree)
	this.RRRotate(tree)
}

func (this *tree_node)getHeight()int{
	lh,rh := 0,0
	if this.leftChild != nil {
		lh = this.leftChild.getLeftHeight() + 1
	}
	if this.rightChild != nil {
		rh = this.rightChild.getRightHeight() + 1
	}
	return Max(lh,rh) + 1
}

func (this *tree_node)getLeftHeight()int{
	if this.leftChild == nil{
		return 0
	}
	return this.leftChild.getHeight()
}

func (this *tree_node)getRightHeight()int{
	if this.rightChild == nil {
		return 0
	}
	return this.rightChild.getHeight()
}

func (this *tree_node)visit(){
	if this.leftChild != nil{
		this.leftChild.visit()
	}

	fmt.Printf("data: %d ",this.data)

	if this.rightChild != nil{
		this.rightChild.visit()
	}
}

type PNode struct{
	Indent	int
	Data	int
}

func pow2(num int)int{
	a := 1
	for i := 1;i <= num;i++{
		a *= 2
	}
	return a
}

func (this *AVLTree)PrintTree()map[int][]PNode{
	printMap := make(map[int][]PNode)
	if this.root == nil{
		fmt.Printf("empty\n")
	}
	num := this.root.getHeight()
	max := pow2(num)
	for i := 0 ; i < num ;i++{
		printMap[i] = make([]PNode,0)
	}
	
	this.root.printNode(0,0,max*4,printMap)
	return printMap
}

func (node *tree_node)printNode(lay,prefix,space int,m map[int][]PNode){
	
	shared := pow2(lay) + 1
	part :=space / shared
	
	n := PNode{Data:node.data,Indent:part}
	m[lay] = append(m[lay],n)
	
	if node.leftChild != nil{
		node.leftChild.printNode(lay+1,prefix,space,m)
	}else{
		_,isExist := m[lay+1]
		if isExist {
			shared := pow2(lay+1) + 1
			part :=space / shared
			n := PNode{Data:-1,Indent:part}
			m[lay+1] = append(m[lay+1],n)
		}
	}

	if node.rightChild != nil{
		node.rightChild.printNode(lay+1,prefix,space,m)
	}else{
		_,isExist := m[lay+1]
		if isExist {
			shared := pow2(lay+1) + 1
			part :=space / shared
			n := PNode{Data:-1,Indent:part}
			m[lay+1] = append(m[lay+1],n)
		}
	}
}